# ------------------------------------------------------------------------------
# Copyright (c) HQU
# Licensed under the HQU License.
# Written by Wang Youjije (youjieWang@stu.hqu.edu.cn)
# Modified by Wang Youjije (youjieWang@stu.hqu.edu.cn)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import logging
import time
from pathlib import Path

import torch
import torch.optim as optim
import torch.nn as nn


def create_logger(cfg, cfg_name, phase='train'):
    '''
    创建运行代码的日志
    :param cfg:参数
    :param cfg_name:cfg的文件名
    :param phase:  train or test
    :return: logger，文件输出路径，tensorboard日志输出的路径
    '''
    root_output_dir = Path(cfg.OUTPUT_DIR)  # 可以通过把字符串类型的路径，转换为 Pathlib.Path 类型的路径
    # set up logger
    # 如果文件路径不存在，就创建文件路径
    if not root_output_dir.exists():
        print('=> creating root_output_dir: {}'.format(root_output_dir))
        root_output_dir.mkdir()  # 创建存放文件的根路径

    # TODO 不知道HYBRID_JOINTS_TYPE这个是什么意思
    # 如果有HYBRID_JOINTS_TYPE， dataset就是就增加一个HYBRID_JOINTS_TYPE，否则就是DATASET
    # HYBRID_JOINTS_TYPE这个好像配置文件中没有这个
    dataset = cfg.DATASET.DATASET + '_' + cfg.DATASET.HYBRID_JOINTS_TYPE \
        if cfg.DATASET.HYBRID_JOINTS_TYPE else cfg.DATASET.DATASET
    dataset = dataset.replace(':', '_')  # 这个就是数据集的路径
    model = cfg.MODEL.NAME  # 模型的名称
    # TP_H_w48_256x192_stage3_1_4_d192_h384_relu_enc4_mh1.yaml--》TP_H_w48_256x192_stage3_1_4_d192_h384_relu_enc4_mh1
    cfg_name = os.path.basename(cfg_name).split('.')[0]  # 除了后缀以外的文件名

    final_output_dir = root_output_dir / dataset / model / cfg_name  # 获取存放模型的整个最终路径

    print('=> creating final_output_dir: {}'.format(final_output_dir))
    final_output_dir.mkdir(parents=True, exist_ok=True)
    time_str = time.strftime('%Y-%m-%d-%H-%M')  # 时间的输出格式,以及获取当前时间
    log_file = '{}_{}_{}.log'.format(cfg_name, time_str, phase)  # 日志的文件名
    final_log_file = final_output_dir / log_file  # 最终日志的存放路径

    head = '%(asctime)-15s %(message)s'
    logging.basicConfig(filename=str(final_log_file),
                        format=head)
    logger = logging.getLogger()  # 初始化logger：logging模块中最基础的对象
    logger.setLevel(logging.INFO)  # logging.INFO确认一切按预期运行

    console = logging.StreamHandler()  # 日志输出的地方，这边设置日志输出到流
    logging.getLogger('').addHandler(console)

    tensorboard_log_dir = Path(cfg.LOG_DIR) / dataset / model / \
                          (cfg_name + '_' + time_str)
    print('=> creating tensorboard_log_dir: {}'.format(tensorboard_log_dir))
    tensorboard_log_dir.mkdir(parents=True, exist_ok=True)
    return logger, str(final_output_dir), str(tensorboard_log_dir)


# 根据定义的参数，使用不同的优化器
def get_optimizer(cfg, model):
    optimizer = None
    if cfg.TRAIN.OPTIMIZER == 'sgd':
        optimizer = optim.SGD(
            model.parameters(),
            lr=cfg.TRAIN.LR,
            momentum=cfg.TRAIN.MOMENTUM,
            weight_decay=cfg.TRAIN.WD,
            nesterov=cfg.TRAIN.NESTEROV
        )
    elif cfg.TRAIN.OPTIMIZER == 'adam':
        optimizer = optim.Adam(
            model.parameters(),
            lr=cfg.TRAIN.LR
        )

    return optimizer


# 保存中间最好的模型
# TODO 4/25这边可能在华为云上面需要修改保存的路径
def save_checkpoint(states, is_best, output_dir,
                    filename='checkpoint.pth'):
    torch.save(states, os.path.join(output_dir, filename))
    if is_best and 'state_dict' in states:
        torch.save(states['best_state_dict'],
                   os.path.join(output_dir, 'model_best.pth'))



# 2/24 增加的代码


